package io.youngledo.usb.license;

import lombok.extern.slf4j.Slf4j;

import java.nio.file.Files;
import java.nio.file.Path;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/// RSA密钥对生成和管理工具
///
/// 用于生成、保存和加载RSA公钥/私钥对。
/// - **私钥**: 仅用于生成许可证（开发者保密）
/// - **公钥**: 内置在应用程序中，用于验证许可证
///
/// ## 使用流程
/// 1. 开发者首次运行 `generateAndSaveKeyPair()` 生成密钥对
/// 2. 私钥保存到安全位置（不提交到Git）
/// 3. 公钥内置到 `LicenseValidator` 中
/// 4. 使用私钥签名许可证，用户端用公钥验证
@Slf4j
public class KeyManager {
    private static final int KEY_SIZE = 2048;
    private static final String ALGORITHM = "RSA";

    /// 生成RSA密钥对
    ///
    /// @return 包含公钥和私钥的KeyPair对象
    /// @throws NoSuchAlgorithmException 如果RSA算法不可用
    public static KeyPair generateKeyPair() throws NoSuchAlgorithmException {
        var keyPairGenerator = KeyPairGenerator.getInstance(ALGORITHM);
        keyPairGenerator.initialize(KEY_SIZE);
        return keyPairGenerator.generateKeyPair();
    }

    /// 生成并保存密钥对到文件
    ///
    /// 私钥和公钥将分别保存为PEM格式文件。
    /// **警告**: 私钥文件必须妥善保管，不要提交到版本控制系统！
    ///
    /// @param privateKeyPath 私钥文件保存路径
    /// @param publicKeyPath 公钥文件保存路径
    /// @throws Exception 如果密钥生成或保存失败
    public static void generateAndSaveKeyPair(Path privateKeyPath, Path publicKeyPath) throws Exception {
        log.info("正在生成RSA密钥对...");
        var keyPair = generateKeyPair();

        // 保存私钥
        savePrivateKey(keyPair.getPrivate(), privateKeyPath);
        log.info("私钥已保存到: {}", privateKeyPath.toAbsolutePath());

        // 保存公钥
        savePublicKey(keyPair.getPublic(), publicKeyPath);
        log.info("公钥已保存到: {}", publicKeyPath.toAbsolutePath());

        log.warn("⚠️  请妥善保管私钥文件，不要提交到版本控制系统！");
        log.info("✅ 请将公钥内容复制到 LicenseValidator.PUBLIC_KEY 常量中");
    }

    /// 保存私钥到文件（PEM格式）
    ///
    /// @param privateKey 私钥对象
    /// @param path 保存路径
    /// @throws Exception 如果保存失败
    public static void savePrivateKey(PrivateKey privateKey, Path path) throws Exception {
        var encoded = Base64.getEncoder().encodeToString(privateKey.getEncoded());
        var pem = """
                -----BEGIN PRIVATE KEY-----
                %s
                -----END PRIVATE KEY-----
                """.formatted(formatBase64(encoded));
        Files.writeString(path, pem);
    }

    /// 保存公钥到文件（PEM格式）
    ///
    /// @param publicKey 公钥对象
    /// @param path 保存路径
    /// @throws Exception 如果保存失败
    public static void savePublicKey(PublicKey publicKey, Path path) throws Exception {
        var encoded = Base64.getEncoder().encodeToString(publicKey.getEncoded());
        var pem = """
                -----BEGIN PUBLIC KEY-----
                %s
                -----END PUBLIC KEY-----
                """.formatted(formatBase64(encoded));
        Files.writeString(path, pem);
    }

    /// 从文件加载私钥
    ///
    /// @param path 私钥文件路径
    /// @return PrivateKey对象
    /// @throws Exception 如果加载或解析失败
    public static PrivateKey loadPrivateKey(Path path) throws Exception {
        var pem = Files.readString(path);
        var privateKeyPEM = pem
                .replace("-----BEGIN PRIVATE KEY-----", "")
                .replace("-----END PRIVATE KEY-----", "")
                .replaceAll("\\s", "");

        var encoded = Base64.getDecoder().decode(privateKeyPEM);
        var keySpec = new PKCS8EncodedKeySpec(encoded);
        var keyFactory = KeyFactory.getInstance(ALGORITHM);
        return keyFactory.generatePrivate(keySpec);
    }

    /// 从文件加载公钥
    ///
    /// @param path 公钥文件路径
    /// @return PublicKey对象
    /// @throws Exception 如果加载或解析失败
    public static PublicKey loadPublicKey(Path path) throws Exception {
        var pem = Files.readString(path);
        var publicKeyPEM = pem
                .replace("-----BEGIN PUBLIC KEY-----", "")
                .replace("-----END PUBLIC KEY-----", "")
                .replaceAll("\\s", "");

        var encoded = Base64.getDecoder().decode(publicKeyPEM);
        var keySpec = new X509EncodedKeySpec(encoded);
        var keyFactory = KeyFactory.getInstance(ALGORITHM);
        return keyFactory.generatePublic(keySpec);
    }

    /// 从PEM字符串加载公钥
    ///
    /// 用于从硬编码的公钥字符串中加载PublicKey对象。
    /// 这个方法将在 LicenseValidator 中使用。
    ///
    /// @param pem PEM格式的公钥字符串
    /// @return PublicKey对象
    /// @throws Exception 如果解析失败
    public static PublicKey loadPublicKeyFromString(String pem) throws Exception {
        var publicKeyPEM = pem
                .replace("-----BEGIN PUBLIC KEY-----", "")
                .replace("-----END PUBLIC KEY-----", "")
                .replaceAll("\\s", "");

        var encoded = Base64.getDecoder().decode(publicKeyPEM);
        var keySpec = new X509EncodedKeySpec(encoded);
        var keyFactory = KeyFactory.getInstance(ALGORITHM);
        return keyFactory.generatePublic(keySpec);
    }

    /// 格式化Base64字符串（每64字符换行）
    ///
    /// @param base64 原始Base64字符串
    /// @return 格式化后的字符串
    private static String formatBase64(String base64) {
        var formatted = new StringBuilder();
        for (int i = 0; i < base64.length(); i += 64) {
            formatted.append(base64, i, Math.min(i + 64, base64.length())).append("\n");
        }
        return formatted.toString();
    }

    /// 主方法：用于生成密钥对
    ///
    /// 运行此方法将在当前目录生成 `private_key.pem` 和 `public_key.pem`
    ///
    /// **使用方法**:
    /// ```bash
    /// java io.youngledo.usb.license.KeyManager
    /// ```
    public static void main(String[] args) {
        try {
            var privateKeyPath = Path.of("private_key.pem");
            var publicKeyPath = Path.of("public_key.pem");

            generateAndSaveKeyPair(privateKeyPath, publicKeyPath);

            System.out.println("\n========================================");
            System.out.println("🔑 RSA密钥对生成成功！");
            System.out.println("========================================");
            System.out.println("私钥文件: " + privateKeyPath.toAbsolutePath());
            System.out.println("公钥文件: " + publicKeyPath.toAbsolutePath());
            System.out.println("\n⚠️  请立即执行以下操作：");
            System.out.println("1. 将 private_key.pem 移动到安全位置");
            System.out.println("2. 添加 private_key.pem 到 .gitignore");
            System.out.println("3. 复制 public_key.pem 的内容到 LicenseValidator.PUBLIC_KEY");
            System.out.println("========================================\n");

            // 输出公钥内容供复制
            System.out.println("📋 公钥内容（复制到 LicenseValidator.PUBLIC_KEY）：\n");
            System.out.println(Files.readString(publicKeyPath));

        } catch (Exception e) {
            log.error("密钥生成失败", e);
            System.err.println("❌ 密钥生成失败: " + e.getMessage());
        }
    }
}
